package org.demos.springaiargdemo1.controller;

import com.alibaba.fastjson2.JSON;
import org.demos.springaiargdemo1.model.dto.*;
import org.demos.springaiargdemo1.service.RAGService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.time.Instant;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

@RestController
@RequestMapping("/v1")
public class OpenAIController {

    private final ChatClient chatClient;
    private final RAGService ragService;

    public OpenAIController(ChatClient.Builder builder, RAGService ragService) {
        this.chatClient = builder.build();
        this.ragService = ragService;
    }


    /**
     * OpenAI 兼容的流式聊天接口
     */
    @PostMapping(value = "/chat/completions/stream", produces = "text/event-stream")
    public Flux<String> streamChatCompletion(
            @RequestBody ChatCompletionRequest request) {

        // 1. 构建用户输入
        String userInput = extractUserInput(request.getMessages());

        // 2. 根据模型选择处理方式
        if (request.getModel().contains("rag")) {
            // RAG 模式流式响应 - 包含文档信息
            AtomicReference<List<RAGResponse.RetrievedDocument>> documentsRef = new AtomicReference<>();
            
            return ragService.generateResponse(userInput)
                    .map(ragResponse -> {
                        // 保存文档信息（只在第一次保存）
                        if (documentsRef.get() == null && ragResponse.getDocuments() != null && !ragResponse.getDocuments().isEmpty()) {
                            documentsRef.set(ragResponse.getDocuments());
                        }
                        
                        return createSSEEventWithDocuments(request.getModel(), ragResponse.getContent(), documentsRef.get());
                    });
        } else {
            // 普通聊天流式响应
            return chatClient.prompt()
                    .user(userInput)
                    .stream()
                    .content()
                    .map(content -> createSSEEvent(request.getModel(), content));
        }
    }

    /**
     * OpenAI 兼容的模型列表接口
     */
    @GetMapping("/models")
    public ResponseEntity<ModelListResponse> listModels() {
        List<ModelInfo> models = new ArrayList<>();

        // 添加普通聊天模型
        models.add(createModelInfo("demo-chat", 16385));

        // 添加 RAG 模型
        models.add(createModelInfo("demo-rag", 128000));

        ModelListResponse response = ModelListResponse.builder()
                .object("list")
                .data(models)
                .build();

        return ResponseEntity.ok(response);
    }

    // 辅助方法：从消息列表中提取用户输入
    private String extractUserInput(List<ChatMessage> messages) {
        StringBuilder userInput = new StringBuilder();

        for (ChatMessage message : messages) {
            if ("user".equals(message.getRole())) {
                if (!userInput.isEmpty()) userInput.append("\n");
                userInput.append(message.getContent());
            }
        }

        return userInput.toString();
    }

    // 辅助方法：创建完成的响应
    private ChatCompletionResponse createCompletionResponse(String model, String content) {
        ChatMessage message = new ChatMessage("assistant", content, null);
        
        ChatChoice choice = ChatChoice.builder()
                .index(0)
                .message(message)
                .finishReason("stop")
                .build();

        Usage usage = Usage.builder()
                .promptTokens(content.length() / 4) // 估算值
                .completionTokens(content.length() / 4) // 估算值
                .totalTokens(content.length() / 2) // 估算值
                .build();

        return ChatCompletionResponse.builder()
                .id("chatcmpl-" + UUID.randomUUID().toString())
                .object("chat.completion")
                .created(Instant.now().getEpochSecond())
                .model(model)
                .choices(List.of(choice))
                .usage(usage)
                .build();
    }

    // 辅助方法：创建 SSE 事件
    private String createSSEEvent(String model, String content) {
        ChatMessage delta = new ChatMessage(null, content, null);
        
        ChatChoice choice = ChatChoice.builder()
                .index(0)
                .delta(delta)
                .finishReason(null)
                .build();

        StreamEvent event = StreamEvent.builder()
                .id("chatcmpl-" + UUID.randomUUID().toString())
                .object("chat.completion.chunk")
                .created(Instant.now().getEpochSecond())
                .model(model)
                .choices(List.of(choice))
                .build();

        return JSON.toJSONString(event);
    }

    // 辅助方法：创建包含文档信息的 SSE 事件
    private String createSSEEventWithDocuments(String model, String content, List<RAGResponse.RetrievedDocument> documents) {
        ChatMessage delta = new ChatMessage(null, content, null);
        
        ChatChoice choice = ChatChoice.builder()
                .index(0)
                .delta(delta)
                .finishReason(null)
                .build();

        // 创建扩展的流事件，包含文档信息
        Map<String, Object> eventData = new HashMap<>();
        eventData.put("id", "chatcmpl-" + UUID.randomUUID().toString());
        eventData.put("object", "chat.completion.chunk");
        eventData.put("created", Instant.now().getEpochSecond());
        eventData.put("model", model);
        eventData.put("choices", List.of(choice));
        
        // 添加文档信息（只在有文档时添加）
        if (documents != null && !documents.isEmpty()) {
            eventData.put("documents", documents);
        }

        return JSON.toJSONString(eventData);
    }

    // 辅助方法：创建模型信息
    private ModelInfo createModelInfo(String modelId, int contextLength) {
        ModelPermission permission = ModelPermission.builder()
                .id("modelperm-" + UUID.randomUUID().toString())
                .object("model_permission")
                .created(1700000000L)
                .allowCreateEngine(false)
                .allowSampling(true)
                .allowLogprobs(true)
                .allowSearchIndices(false)
                .allowView(true)
                .allowFineTuning(false)
                .organization("*")
                .group(null)
                .isBlocking(false)
                .build();

        ModelMetadata metadata = ModelMetadata.builder()
                .contextLength(contextLength)
                .capabilities(List.of("chat", "function_calling", "rag"))
                .ragSupport(modelId.contains("rag"))
                .build();

        return ModelInfo.builder()
                .id(modelId)
                .object("model")
                .created(1700000000L)
                .ownedBy("your-organization")
                .permission(List.of(permission))
                .root(modelId)
                .parent(null)
                .metadata(metadata)
                .build();
    }

    // 辅助方法：错误响应
    private ErrorResponse createErrorResponse(String message) {
        ErrorResponse.ErrorDetail errorDetail = ErrorResponse.ErrorDetail.builder()
                .message(message)
                .type("invalid_request_error")
                .code("invalid_model")
                .build();

        return ErrorResponse.builder()
                .error(errorDetail)
                .build();
    }
}
